package com.lian.aieduc.controller;

import com.lian.aieduc.utils.Result;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.Parameter;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.zhipuai.ZhiPuAiChatModel;
import org.springframework.ai.zhipuai.ZhiPuAiChatOptions;
import org.springframework.ai.zhipuai.api.ZhiPuAiApi;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.Objects;

@RestController
@Slf4j
@RequestMapping("/chat")
@RequiredArgsConstructor
@Tag(name = "聊天AI模型相关接口", description = "聊天AI模型相关接口")
public class ChatAIController {
    private final ZhiPuAiChatModel chatModel;

    @Operation(summary = "流式获取AI消息")
    @Parameter(description = "问题")
    @GetMapping(value = "/ai", params = "question", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> getZhiPuAIMessage(String question) {
        Prompt prompt = new Prompt(question, ZhiPuAiChatOptions.builder().withModel(
            ZhiPuAiApi.ChatModel.GLM_4_Flash.getValue()).withTemperature(0.9).build());

        // 流式返回
        return chatModel.stream(prompt).filter(Objects::nonNull).filter(
                chatResponse -> chatResponse.getResults() != null).flatMap(
                chatResponse -> Flux.fromIterable(chatResponse.getResults())).filter(Objects::nonNull).map(
                Generation::getOutput).filter(Objects::nonNull).filter(content -> Objects.nonNull(content.getContent()))
            .map(AssistantMessage::getContent).filter(Objects::nonNull).map(
                content -> ServerSentEvent.builder(content).build()).doOnNext(System.out::println).concatWith(
                Flux.just(ServerSentEvent.builder("complete").build()));
    }


}